#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen);   // 将 main 函数翻译为汇编代码
void translate_callee(CodeGen *codegen); // 将 callee 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_callee(codegen);
    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// 补全 translate_callee 函数
void translate_callee(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明 callee 函数 */
    codegen->append_inst(".globl callee", ASMInstruction::Atrribute);
    codegen->append_inst(".type callee, @function", ASMInstruction::Atrribute);

    /* callee 函数开始 */
    codegen->append_inst("callee", ASMInstruction::Label); // main 函数标签

    /* callee 函数的 Prologue (序言) */
    // 开辟栈帧空间，使用 32 字节
    codegen->append_inst("addi sp, sp, -32");
    // 保存返回地址
    codegen->append_inst("sd ra, 24(sp)");
    // 保存老的 fp
    codegen->append_inst("sd s0, 16(sp)");
    // 设置新的 fp
    codegen->append_inst("addi s0, sp, 32"); 

    /* callee 函数的 label_entry */
    codegen->append_inst(".callee_label_entry", ASMInstruction::Label);

    /* %op1 = alloca i32 */
    codegen->append_inst("%op1 = alloca i32", ASMInstruction::Comment);
    // 为 %op1 分配栈上空间，使用正偏移量
    offset_map["%op1"] = 8; 
    // 将栈上地址存入 %op1
    codegen->append_inst("addi", {"t0", "sp", std::to_string(offset_map["%op1"])});
    codegen->append_inst("sd", {"t0", "0(sp)"});

    /* store i32 %arg0, i32* %op1 */
    codegen->append_inst("store i32 %arg0, i32* %op1", ASMInstruction::Comment);
    // 从 a0 寄存器获取参数值
    codegen->append_inst("ld", {"t0", "0(sp)"}); // 加载 %op1 的地址
    codegen->append_inst("sw", {"a0", "0(t0)"}); // 存储参数值到 %op1

    /* %op2 = load i32, i32* %op1 */
    codegen->append_inst("%op2 = load i32, i32* %op1", ASMInstruction::Comment);
    offset_map["%op2"] = 4; 
    codegen->append_inst("ld", {"t0", "0(sp)"}); // 加载 %op1 的地址
    codegen->append_inst("lw", {"t1", "0(t0)"}); // 加载 %op1 的值
    codegen->append_inst("sw", {"t1", std::to_string(offset_map["%op2"])+"(sp)"});

    /* %op3 = mul i32 3, %op2 */
    codegen->append_inst("%op3 = mul i32 3, %op2", ASMInstruction::Comment);
    offset_map["%op3"] = 0; 
    codegen->append_inst("lw", {"t0", std::to_string(offset_map["%op2"])+"(sp)"});
    codegen->append_inst("li", {"t1", "3"});
    codegen->append_inst("mul", {"t0", "t0", "t1"});
    codegen->append_inst("sw", {"t0", std::to_string(offset_map["%op3"])+"(sp)"});

    /* ret i32 %op3 */
    codegen->append_inst("ret i32 %op3", ASMInstruction::Comment);
    // 将 %op3 的值写入 a0 寄存器
    codegen->append_inst("lw", {"a0", std::to_string(offset_map["%op3"])+"(sp)"});
    codegen->append_inst("j callee_exit");

    /* callee 函数的 Epilogue (收尾) */
    codegen->append_inst("callee_exit", ASMInstruction::Label);
    // 恢复 ra
    codegen->append_inst("ld ra, 24(sp)");
    // 恢复 fp
    codegen->append_inst("ld s0, 16(sp)");
    // 释放栈帧空间
    codegen->append_inst("addi sp, sp, 32");
    // 返回
    codegen->append_inst("jr ra");
}

// 补全 translate_main 函数
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明 main 函数 */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label); // main 函数标签

    /* main 函数的 Prologue (序言) */
    // 开辟栈帧空间
    codegen->append_inst("addi sp, sp, -16");
    // 保存返回地址
    codegen->append_inst("sd ra, 8(sp)");
    // 保存老的 fp
    codegen->append_inst("sd s0, 0(sp)");
    // 设置新的 fp
    codegen->append_inst("addi s0, sp, 16"); 

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);

    /* %op0 = call i32 @callee(i32 110) */
    codegen->append_inst("%op0 = call i32 @callee(i32 110)",
                         ASMInstruction::Comment);
    offset_map["%op0"] = 4; 
    // 将参数写入 a0 寄存器
    codegen->append_inst("li", {"a0", "110"});
    // 调用 callee 函数
    codegen->append_inst("call", {"callee"});
    // 保存返回值
    codegen->append_inst("sw", {"a0", std::to_string(offset_map["%op0"])+"(sp)"});

    /* ret i32 %op0 */
    codegen->append_inst("ret i32 %op0", ASMInstruction::Comment);
    codegen->append_inst("lw", {"a0", std::to_string(offset_map["%op0"])+"(sp)"});
    codegen->append_inst("j main_exit");

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    // 恢复 ra
    codegen->append_inst("ld ra, 8(sp)");
    // 恢复 fp
    codegen->append_inst("ld s0, 0(sp)");
    // 释放栈帧空间
    codegen->append_inst("addi sp, sp, 16");
    // 返回
    codegen->append_inst("jr ra");
}